import cv2

import tracker
from detector import Detector


def function2(vid_path = './data/test1.mp4'):
    # 初始化 yolov5
    detector = Detector()
    # 打开视频
    capture = cv2.VideoCapture(vid_path)

    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    frame_width = int(capture.get(3))
    frame_height = int(capture.get(4))
    out = cv2.VideoWriter('output1.avi', fourcc, 20.0, (frame_width, frame_height))
    cls_id = None
    pos_id = None
    while True:
        # 读取每帧图片
        _, im = capture.read()
        if im is None:
            break

        # # 缩小尺寸，1920x1080->960x540
        # im = cv2.resize(im, (960, 540))

        list_bboxs = []
        bboxes = detector.detect(im)

        # 如果画面中 有bbox
        if len(bboxes) > 0:
            # 接受目标检测结果进行更新
            list_bboxs = tracker.update(bboxes, im)
            # 预测
            # kf = KalmanFilter()
            # for _,_,_,_,_,_,center_x, center_y in list_bboxs:
            #     center_x = int(center_x)
            #     center_y = int(center_y)
            #     predicted = kf.predict([center_x, center_y])
            #     print(predicted)


            # 画框
            # 只有经过deepsort算法验证的的boxs，才会画出来。
            output_image_frame, cls_id, pos_id = tracker.draw_bboxes(im, list_bboxs, line_thickness=None)
            pass
        else:
            # 如果画面中 没有bbox
            output_image_frame = im
        pass
        # cv2.putText(im, f"Detection of target vehicle is block, vehicle type is {cls_id}, ID is {pos_id}", (100, 100),
        #             cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        out.write(output_image_frame)

        cv2.imshow('demo', output_image_frame)
        cv2.waitKey(1)

        pass
    pass

    capture.release()
    out.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    function2()